#include <odometry_chassis/odometry_chassis.hpp>

using namespace std::chrono;
constexpr double WHEEL_RADIUS = 0.09;  // 轮半径 m
constexpr double WHEEL_BASE = 0.361;   // 两轮轴距 m

OdometryChassis::OdometryChassis() : Node("odometry_chassis") {
  initParamters();
  if (!initialize()) {
    RCLCPP_ERROR(this->get_logger(), "Robot init failed, shutting down.");
    rclcpp::shutdown();
  }
  initSubscribers();
  initPublishers();
  timer_ =
      this->create_wall_timer(std::chrono::milliseconds(50),
                              std::bind(&OdometryChassis::updateLoop, this));
}

OdometryChassis::~OdometryChassis() { robot.robot_deinit(); }

void OdometryChassis::initParamters() {
  odom.header.frame_id = "odom";
  odom.child_frame_id = "base_link";
  util::param(this, "publish_tf", publish_tf, publish_tf);
  util::param(this, "low_level_control", low_level_control, low_level_control);
  util::param(this, "odom_frame_id", odom.header.frame_id,
              odom.header.frame_id);
  util::param(this, "base_frame_id", odom.child_frame_id, odom.child_frame_id);

  odom.pose.pose.position.x = 0.0;
  odom.pose.pose.position.y = 0.0;
  yaw = 0.0;
}
void OdometryChassis::initSubscribers() {
  cmd_sub_ = this->create_subscription<geometry_msgs::msg::Twist>(
      "cmd_vel", 10,
      std::bind(&OdometryChassis::cmdVelCallback, this, std::placeholders::_1));
}
void OdometryChassis::initPublishers() {
  odom_pub_ = this->create_publisher<nav_msgs::msg::Odometry>("odom", 10);
  tf_broadcaster_ = std::make_shared<tf2_ros::TransformBroadcaster>(*this);
}

bool OdometryChassis::initialize() {
  try {
    if (!robot.robot_connect()) {  // 尝试连接机器人
      RCLCPP_ERROR(this->get_logger(), "Robot connection failed.");
      return false;
    }
    RCLCPP_INFO(this->get_logger(), "Robot connected successfully.");

    //选择高层控制模式
    if (!robot.switchControlMode(MotorControlMode::HIGH_LEVEL)) {
      RCLCPP_ERROR(this->get_logger(), "Failed to switch to HIGH_LEVEL mode.");
      return 1;
    }

    if (!robot.robot_init()) {  // 尝试初始化机器人（动作）
      RCLCPP_ERROR(this->get_logger(), "Robot initialization failed.");
      return false;
    }
    last_cmd_time_ = steady_clock::now();
    RCLCPP_INFO(this->get_logger(), "Robot initialized successfully.");
    return true;
  } catch (std::exception &e) {
    RCLCPP_ERROR(this->get_logger(), "Failed to initialized, %s", e.what());
  }
  return false;
}

void OdometryChassis::cmdVelCallback(
    const geometry_msgs::msg::Twist::SharedPtr msg) {
  last_cmd_time_ = steady_clock::now();
  sendVelocity(msg->linear.x, msg->angular.z);
}

double OdometryChassis::motorToWheelVelocity(double omega_motor_rad_s,
                                             double wheel_radius,
                                             double gear_ratio) {
  // omega_motor_rad_s: 电机角速度 rad/s
  // wheel_radius: 轮半径 m
  // gear_ratio: 减速比 N
  return wheel_radius * omega_motor_rad_s / gear_ratio;
}

void OdometryChassis::chassisVelToMotorSpeed(double v, double w, double &wl,
                                             double &wr) {
  double vl = v - 0.5 * w * WHEEL_BASE;
  double vr = v + 0.5 * w * WHEEL_BASE;
  wl = vl / WHEEL_RADIUS;
  wr = vr / WHEEL_RADIUS;
}
void OdometryChassis::updateLoop() {
  auto now = steady_clock::now();
  double dt = duration_cast<duration<double>>(now - last_odom_time_).count();
  last_odom_time_ = now;
  // 速度指令超时检查
  double cmd_dt = duration_cast<duration<double>>(now - last_cmd_time_).count();
  if (cmd_dt > 1.0) {
    sendVelocity(0.0, 0.0);
  }
  calcOdometry(dt);
  publishOdometry();
}

void OdometryChassis::calcOdometry(double dt) {
  Motor_Information left_state, right_state;
  robot.getMotorState(EtherCAT_Motor_Index::MOTOR_WHEEL_LEFT, left_state);
  robot.getMotorState(EtherCAT_Motor_Index::MOTOR_WHEEL_RIGHT, right_state);
  state_estimator_zcmt imu_state;
  robot.getIMU_State(imu_state);

  double vl =
      motorToWheelVelocity(left_state.Speed_Actual, WHEEL_RADIUS);  // m/s
  double vr = motorToWheelVelocity(right_state.Speed_Actual, WHEEL_RADIUS);

  odom.twist.twist.linear.x = 0.5 * (vr + vl);
  odom.twist.twist.angular.z = (vr - vl) / WHEEL_BASE;

  // double wheel_yaw_rate = (vr - vl) / WHEEL_BASE;
  // double imu_yaw_rate = imu_state.omega[2];  // rad/s
  // double alpha = 0.8;                        // 融合系数
  // odom.twist.twist.angular.z =
  //     alpha * imu_yaw_rate + (1.0 - alpha) * wheel_yaw_rate;

  odom.pose.pose.position.x += odom.twist.twist.linear.x * cos(yaw) * dt;
  odom.pose.pose.position.y += odom.twist.twist.linear.x * sin(yaw) * dt;
  yaw += odom.twist.twist.angular.z * dt;
  odom.pose.pose.orientation.w = cos(yaw / 2.0);
  odom.pose.pose.orientation.z = sin(yaw / 2.0);
}

void OdometryChassis::publishOdometry() {
  odom.header.stamp = this->now();
  odom_pub_->publish(odom);
  if (publish_tf) {  // 只有当参数为 true 时才发布 TF
    geometry_msgs::msg::TransformStamped odom_tf;
    odom_tf.header.stamp = this->now();
    odom_tf.header.frame_id = odom.header.frame_id;
    odom_tf.child_frame_id = odom.child_frame_id;

    odom_tf.transform.translation.x = odom.pose.pose.position.x;
    odom_tf.transform.translation.y = odom.pose.pose.position.y;
    odom_tf.transform.translation.z = 0.0;
    odom_tf.transform.rotation = odom.pose.pose.orientation;
    tf_broadcaster_->sendTransform(odom_tf);
  }
}

void OdometryChassis::sendVelocity(double linear, double angular) {
  if (low_level_control) {
    sendChassisLowCmd(linear, angular);
  } else {
    sendChassisHightCmd(linear, angular);
  }
}

void OdometryChassis::sendChassisHightCmd(double lin, double ang) {
  ChassisMotion linear_motion, angular_motion;
  Motor_Control linear_control, angular_control;
  // 线速度控制
  if (lin >= 0.0) {
    linear_motion = ChassisMotion::FORWARD;
    linear_control.Speed = lin;
  } else {
    linear_motion = ChassisMotion::BACKWARD;
    linear_control.Speed = fabs(lin);
  }

  // 角速度控制
  if (ang >= 0.0) {
    angular_motion = ChassisMotion::LEFT;
    angular_control.Speed = ang;
  } else {
    angular_motion = ChassisMotion::RIGHT;
    angular_control.Speed = fabs(ang);
  }

  robot.setChassis_high(linear_motion, linear_control);
  robot.setChassis_high(angular_motion, angular_control);
}

void OdometryChassis::sendChassisLowCmd(double lin, double ang) {
  double wl, wr;
  chassisVelToMotorSpeed(lin, ang, wl, wr);
  Motor_Control control;
  control.Speed = wl;
  robot.setMotorControl_low(EtherCAT_Motor_Index::MOTOR_WHEEL_LEFT, control);
  control.Speed = wr;
  robot.setMotorControl_low(EtherCAT_Motor_Index::MOTOR_WHEEL_RIGHT, control);
}
